{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.document_loaders import TextLoader\n",
    "loader = TextLoader('./state_of_the_union.txt', encoding='utf8')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.embeddings.openai import OpenAIEmbeddings\n",
    "from langchain.text_splitter import CharacterTextSplitter\n",
    "from langchain.vectorstores import Chroma\n",
    "\n",
    "with open('./state_of_the_union.txt') as f:\n",
    "    state_of_the_union = f.read()\n",
    "text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)\n",
    "texts = text_splitter.split_text(state_of_the_union)\n",
    "\n",
    "# embeddings = OpenAIEmbeddings()\n",
    "# docsearch = Chroma.from_texts(texts, embeddings)\n",
    "\n",
    "# query = \"What did the president say about Ketanji Brown Jackson\"\n",
    "# docs = docsearch.similarity_search(query)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "d:\\Anaconda\\envs\\pytorch\\lib\\site-packages\\scipy\\__init__.py:146: UserWarning: A NumPy version >=1.17.3 and <1.25.0 is required for this version of SciPy (detected version 1.26.2\n",
      "  warnings.warn(f\"A NumPy version >={np_minversion} and <{np_maxversion}\"\n"
     ]
    }
   ],
   "source": [
    "from typing import Any\n",
    "from sentence_transformers import SentenceTransformer\n",
    "from chromadb import Documents, EmbeddingFunction, Embeddings\n",
    "\n",
    "class Embeddings(EmbeddingFunction):\n",
    "    def __init__(self, model_path):\n",
    "        self.model = SentenceTransformer(model_path)\n",
    "\n",
    "    def __call__(self, input: Documents) -> Embeddings:\n",
    "        # embed the documents somehow\n",
    "        embeddings = [self.embed_query(x) for x in input]\n",
    "        \n",
    "        # return embeddings\n",
    "    # def embed_documents(self, input):\n",
    "    #     embeddings = self.model.encode(input)\n",
    "    #     encod_list = embeddings.tolist()\n",
    "    #     return encod_list\n",
    "    def embed_documents(self, input):\n",
    "        embeddings = [self.embed_query(x) for x in input]\n",
    "        # embeddings = self.model.encode(input)\n",
    "        # encod_list = embeddings.tolist()\n",
    "        # return encod_list\n",
    "        return embeddings\n",
    "\n",
    "    def embed_query(self, input):\n",
    "        embeddings = self.model.encode([input])\n",
    "        encod_list = embeddings.tolist()\n",
    "        return encod_list[0]\n",
    "# mymodel = Embeddings(model_config.embedding_model)\n",
    "# class MyFunction(EmbeddingFunction):\n",
    "#     def __call__(self, texts: Documents) -> Embeddings:\n",
    "#         embeddings = [mymodel.embed_query(x) for x in texts]\n",
    "#         return embeddings\n",
    "#     def embed_documents(self, text_list):\n",
    "#         embeddings = mymodel.encode(text_list)\n",
    "#         encod_list = embeddings.tolist()\n",
    "#         return encod_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from chromadb import Documents, EmbeddingFunction, Embeddings\n",
    "\n",
    "class Mymodel(EmbeddingFunction):\n",
    "    def __call__(self, texts: Documents) -> Embeddings:\n",
    "        # embed the documents somehow\n",
    "        return [[1, 2, 3], [2, 4, 6]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No sentence-transformers model found with name ./model/m3e. Creating a new one with MEAN pooling.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4\n",
      "4\n"
     ]
    }
   ],
   "source": [
    "from langchain.embeddings.openai import OpenAIEmbeddings\n",
    "from langchain.text_splitter import CharacterTextSplitter\n",
    "from langchain.vectorstores import Chroma\n",
    "from embedding import Embeddings\n",
    "from configs.params import ModelParams\n",
    "model_config = ModelParams()\n",
    "with open('prompt.txt') as f:\n",
    "    state_of_the_union = f.read()\n",
    "text_splitter = CharacterTextSplitter(chunk_size=500, chunk_overlap=0)\n",
    "texts = text_splitter.split_text(state_of_the_union)\n",
    "\n",
    "\n",
    "embeddings = Embeddings(model_config.embedding_model)\n",
    "metadatas = [{\"title\":'prompt2.txt', 'segment': i} for i in range(1, len(texts) + 1)]\n",
    "ids = [str(i) for i in range(1, len(texts) + 1)] \n",
    "docsearch = Chroma.from_texts(texts, embeddings, metadatas=metadatas)\n",
    "\n",
    "query = \"厂妹\"\n",
    "# query = \"What did the president say about Ketanji Brown Jackson\"\n",
    "docs = docsearch.similarity_search(query)\n",
    "print(len(docs))\n",
    "print(docsearch._collection.count())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res = docsearch.get(where={'title': 'prompt2.txt'})\n",
    "res2 = docsearch.get(where={'title': 'prompt.txt'})\n",
    "len(res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(res2['ids'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(page_content='今天早上，我参加了工程伦理领读课程，由几位报名的小组同学进行内容分享。他们的分享让我对《素质教育与前沿技术》这门课程有了更深刻的理解。课程内容涵盖了六个主题，每个主题都深入探讨了工程伦理的不同方面。通过分析真实的案例，我认识到伦理决策的重要性，并学会了如何从伦理的角度分析问题，以便做出符合道德标准的决策。同时，我也了解到实践伦理工具箱的重要性，它为工程师提供了一系列方法和工具来应对伦理挑战。这些工具不仅有助于解决问题，还能帮助工程师更好地理解和管理伦理风险。此外，课程还强调了工程师的责任，包括对他们的工作承担道德和法律责任。这一主题教导我们如何在工程实践中积极履行职业道德。同时，我也认识到工程师不仅是技术专家，还需要考虑他们的工作对社会和价值观的影响。建立信任关系在工程领域中的重要性也得到了强调，工程师应该努力提供可靠的解决方案，以维护客户和社会对他们的信任。最后，课程还探讨了工程风险与责任的问题，强调了工程师在面对潜在风险时应该如何行动。总之，这门课程使我深刻地认识到工程伦理对于工程师的重要性。工程师不仅是技术专家，还是社会的一部分，他们的工作对社会和环境有深远的影响。因此，了解和积极践行职业道德是至关重要的。', metadata={'source': './data/case.txt'})]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.text_splitter import CharacterTextSplitter\n",
    "from langchain.document_loaders import TextLoader\n",
    "def load_file(filepath, chunk_size, chunk_overlap):\n",
    "    loader = TextLoader(filepath, encoding='utf-8')\n",
    "    documents = loader.load()\n",
    "    text_splitter = CharacterTextSplitter(separator='\\n', chunk_size=chunk_size, chunk_overlap=chunk_overlap)\n",
    "    docs = text_splitter.split_documents(documents)\n",
    "    return docs\n",
    "docs = load_file('./data/case.txt', 300, 0)\n",
    "docs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "docsearch._collection.delete(where={'title': 'prompt.txt'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'page_content': '彼时，大佬们的主场还是各大独角兽，谷歌是大部分码工的终极理想，亚麻还是一个名声不佳的血汗工厂。找工作的这几个月，我从朋友圈里，欣赏了一号公路的美景，了解了湾区各个中餐厅的味道，领略了各个大厂的食堂、健身房、搬家福利。加州的阳光，照耀着这些天之骄子。他们去冲浪，去潜水，去跳伞，他们站在浪潮之巅，意气风发。还没有春风得意，提前看了不少长安花。以至于终于拿了一个Offer后，我的心中毫无波澜。久未联系的同学，约我爬mission peak，在光秃秃的山顶，同学得知我的包裹，微微一笑，说，不错，在我们班同学里能排前列——当然略逊于他。进厂后，厂妹的工作没什么技术含量，也没什么前途，但借着硅谷的光，也在风口上飘了好一会儿。我听同事讲她在爱马仕配货的故事，也咬咬牙，在奥特莱斯买了一个LV。那段时间，还认识了很多车标。美国是车轮上的国家，湾区更是如此。第一次坐特斯拉是打顺风车，车门都不会开。那会儿，看到刚入职的新人，拿了签字费就纷纷开上宝马、奔驰。这让车盲的我，一度以为它们和现代一样，是买菜代步车。 这一年，厂妹和打工仔在厂里的食堂相识了。和工人文学里老一辈的爱情故事一样，无非就是公司群活动认识了，一来二去地约饭熟悉了，心照不宣地牵手了。这一年的H1B抽签，硕士毕业生号称有史高的中签率。果然，朋友圈全部中签，除了厂妹和打工仔。自此，H1B成为悬在头上的达摩克利斯之剑。不久后，疫情封城，H1B二抽没中，股票三次熔断。 因为疫情，朋友逐渐失联，恋情被迫稳定。 21年初，股票和就业市场一下就触底反弹，高歌猛进。失联的朋友又逐渐聚起来，大家在湾区这片遍地黄金的土地，也已经腌臜入味。饭局里充斥着 “你是大佬”“你才是大佬”的极限拉扯，各自update自己及朋友的包裹、职级、福利，依次聊一遍股市、房价、虚拟币、新能源、自动驾驶。 湾区的阳光格外刺眼，照耀着幸运的弄潮儿们。 三年之期到，大家磨刀霍霍向房市。更多的人，开始跳槽，奔赴几十万美元的大包。这是最好的时代，一切都欣欣向荣，每个人都有光明且富裕的未来。除了只剩最后一次抽签机会的厂妹，和她同样倒霉的另一半。最后一次抽签，没中，意料之中。得知这个消息时，厂妹正和对象驱车驶在俄勒冈州的公路上，遇见罕见的太阳雨，天边出现一朵七彩的云。隐约记起，刚去湾区的时候，也看到了太阳雨，还有双彩虹。 正值疫情，SFO的机场空空如也。 厂妹拖着行李箱，感到前所未有的轻松。回想起湾区的一切，深感，这是一个最好的时代。',\n",
       " 'metadata': {'segment': 2, 'title': 'prompt2.txt'},\n",
       " 'type': 'Document'}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dict(docs[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "彼时，大佬们的主场还是各大独角兽，谷歌是大部分码工的终极理想，亚麻还是一个名声不佳的血汗工厂。找工作的这几个月，我从朋友圈里，欣赏了一号公路的美景，了解了湾区各个中餐厅的味道，领略了各个大厂的食堂、健身房、搬家福利。加州的阳光，照耀着这些天之骄子。他们去冲浪，去潜水，去跳伞，他们站在浪潮之巅，意气风发。还没有春风得意，提前看了不少长安花。以至于终于拿了一个Offer后，我的心中毫无波澜。久未联系的同学，约我爬mission peak，在光秃秃的山顶，同学得知我的包裹，微微一笑，说，不错，在我们班同学里能排前列——当然略逊于他。进厂后，厂妹的工作没什么技术含量，也没什么前途，但借着硅谷的光，也在风口上飘了好一会儿。我听同事讲她在爱马仕配货的故事，也咬咬牙，在奥特莱斯买了一个LV。那段时间，还认识了很多车标。美国是车轮上的国家，湾区更是如此。第一次坐特斯拉是打顺风车，车门都不会开。那会儿，看到刚入职的新人，拿了签字费就纷纷开上宝马、奔驰。这让车盲的我，一度以为它们和现代一样，是买菜代步车。 这一年，厂妹和打工仔在厂里的食堂相识了。和工人文学里老一辈的爱情故事一样，无非就是公司群活动认识了，一来二去地约饭熟悉了，心照不宣地牵手了。这一年的H1B抽签，硕士毕业生号称有史高的中签率。果然，朋友圈全部中签，除了厂妹和打工仔。自此，H1B成为悬在头上的达摩克利斯之剑。不久后，疫情封城，H1B二抽没中，股票三次熔断。 因为疫情，朋友逐渐失联，恋情被迫稳定。 21年初，股票和就业市场一下就触底反弹，高歌猛进。失联的朋友又逐渐聚起来，大家在湾区这片遍地黄金的土地，也已经腌臜入味。饭局里充斥着 “你是大佬”“你才是大佬”的极限拉扯，各自update自己及朋友的包裹、职级、福利，依次聊一遍股市、房价、虚拟币、新能源、自动驾驶。 湾区的阳光格外刺眼，照耀着幸运的弄潮儿们。 三年之期到，大家磨刀霍霍向房市。更多的人，开始跳槽，奔赴几十万美元的大包。这是最好的时代，一切都欣欣向荣，每个人都有光明且富裕的未来。除了只剩最后一次抽签机会的厂妹，和她同样倒霉的另一半。最后一次抽签，没中，意料之中。得知这个消息时，厂妹正和对象驱车驶在俄勒冈州的公路上，遇见罕见的太阳雨，天边出现一朵七彩的云。隐约记起，刚去湾区的时候，也看到了太阳雨，还有双彩虹。 正值疫情，SFO的机场空空如也。 厂妹拖着行李箱，感到前所未有的轻松。回想起湾区的一切，深感，这是一个最好的时代。\n",
      "prompt2.txt\n",
      "2\n",
      "这是一个最好的时代。2018年，我身无offer，没有实习，刷题稀烂，来到传说中的硅谷。没有人叫这里硅谷，大家都亲切地叫湾区。 湾区的阳光格外耀眼。一栋栋方方正正的科技公司大楼，在阳光的映照下，像一块块巨大的金砖。\n",
      "\n",
      "毫无美感，但着实诱人。房东是一对年轻的码农夫妇，对我照顾有加。男房东刚刚考了地产经济执照，让我以后有买房需求找他。\n",
      "\n",
      "我说，我还没找到工作呢。他说，没问题的，他帮我内推。学长听闻我来加州，请我吃火锅。饭桌上，他痛斥狗家只会刷题，工作效率低，亚麻YYDS。\n",
      "\n",
      "随着一盘一盘的肉下锅，学长越发义愤填膺，一会儿骂白左是sb，一会儿感慨北美华人不团结。猛然间，话锋一转，他告诉我，他和他的同学们，都是工作三年买的房。\n",
      "\n",
      "我说，我还没找到工作呢。他说，没问题的，他帮我内推。这是一个最好的时代。虽然我一无所有，但所有人都充满信心。\n",
      "prompt2.txt\n",
      "1\n"
     ]
    }
   ],
   "source": [
    "for doc in docs:\n",
    "    doc = dict(doc)\n",
    "    # print(doc)\n",
    "    print(doc['page_content'])\n",
    "    print(doc['metadata']['title'])\n",
    "    print(doc['metadata']['segment'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# docsearch.delete_collection()\n",
    "docsearch._collection.count()\n",
    "# docs[0].page_content\n",
    "# docs[0].metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Document(page_content='彼时，大佬们的主场还是各大独角兽，谷歌是大部分码工的终极理想，亚麻还是一个名声不佳的血汗工厂。找工作的这几个月，我从朋友圈里，欣赏了一号公路的美景，了解了湾区各个中餐厅的味道，领略了各个大厂的食堂、健身房、搬家福利。加州的阳光，照耀着这些天之骄子。他们去冲浪，去潜水，去跳伞，他们站在浪潮之巅，意气风发。还没有春风得意，提前看了不少长安花。以至于终于拿了一个Offer后，我的心中毫无波澜。久未联系的同学，约我爬mission peak，在光秃秃的山顶，同学得知我的包裹，微微一笑，说，不错，在我们班同学里能排前列——当然略逊于他。进厂后，厂妹的工作没什么技术含量，也没什么前途，但借着硅谷的光，也在风口上飘了好一会儿。我听同事讲她在爱马仕配货的故事，也咬咬牙，在奥特莱斯买了一个LV。那段时间，还认识了很多车标。美国是车轮上的国家，湾区更是如此。第一次坐特斯拉是打顺风车，车门都不会开。那会儿，看到刚入职的新人，拿了签字费就纷纷开上宝马、奔驰。这让车盲的我，一度以为它们和现代一样，是买菜代步车。 这一年，厂妹和打工仔在厂里的食堂相识了。和工人文学里老一辈的爱情故事一样，无非就是公司群活动认识了，一来二去地约饭熟悉了，心照不宣地牵手了。这一年的H1B抽签，硕士毕业生号称有史高的中签率。果然，朋友圈全部中签，除了厂妹和打工仔。自此，H1B成为悬在头上的达摩克利斯之剑。不久后，疫情封城，H1B二抽没中，股票三次熔断。 因为疫情，朋友逐渐失联，恋情被迫稳定。 21年初，股票和就业市场一下就触底反弹，高歌猛进。失联的朋友又逐渐聚起来，大家在湾区这片遍地黄金的土地，也已经腌臜入味。饭局里充斥着 “你是大佬”“你才是大佬”的极限拉扯，各自update自己及朋友的包裹、职级、福利，依次聊一遍股市、房价、虚拟币、新能源、自动驾驶。 湾区的阳光格外刺眼，照耀着幸运的弄潮儿们。 三年之期到，大家磨刀霍霍向房市。更多的人，开始跳槽，奔赴几十万美元的大包。这是最好的时代，一切都欣欣向荣，每个人都有光明且富裕的未来。除了只剩最后一次抽签机会的厂妹，和她同样倒霉的另一半。最后一次抽签，没中，意料之中。得知这个消息时，厂妹正和对象驱车驶在俄勒冈州的公路上，遇见罕见的太阳雨，天边出现一朵七彩的云。隐约记起，刚去湾区的时候，也看到了太阳雨，还有双彩虹。 正值疫情，SFO的机场空空如也。 厂妹拖着行李箱，感到前所未有的轻松。回想起湾区的一切，深感，这是一个最好的时代。', metadata={'segment': 2, 'text': 'prompt.txt'})"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "docs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "100"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "docsearch._collection.count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'我说，我还没找到工作呢。他说，没问题的，他帮我内推。学长听闻我来加州，请我吃火锅。饭桌上，他痛斥狗家只会刷题，工作效率低，亚麻YYDS。\\n\\n随着一盘一盘的肉下锅，学长越发义愤填膺，一会儿骂白左是sb，一会儿感慨北美华人不团结。猛然间，话锋一转，他告诉我，他和他的同学们，都是工作三年买的房。'"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with open('prompt.txt') as f:\n",
    "    state_of_the_union = f.read()\n",
    "text_splitter = CharacterTextSplitter(chunk_size=200, chunk_overlap=0)\n",
    "texts = text_splitter.split_text(state_of_the_union)\n",
    "print(len(texts))\n",
    "texts[0]\n",
    "texts[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No sentence-transformers model found with name ./model/m3e. Creating a new one with MEAN pooling.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1024\n"
     ]
    }
   ],
   "source": [
    "from embedding import Embeddings\n",
    "from configs.params import ModelParams\n",
    "model_config = ModelParams()\n",
    "embeddings = Embeddings(model_config.embedding_model)\n",
    "text = \"hello world\"\n",
    "embed = embeddings.embed_documents(text)\n",
    "print(len(embed))\n",
    "\n",
    "# openai.api_key = \"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\"\n",
    "# PROMPT_TEMPLATE = \"\"\"已知信息：\n",
    "# {context} \n",
    "\n",
    "# 根据上述已知信息，简洁和专业的来回答用户的问题。如果无法从中得到答案，请说 “根据已知信息无法回答该问题” 或 “没有提供足够的相关信息”，不允许在答案中添加编造成分，答案请使用中文。 问题是：{question}\"\"\"\n",
    "\n",
    "\n",
    "# es = ES(model_config.embedding_model)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
